import pickle

from pytdx.hq import TdxHq_API


# reuse api?
class HQ:
    def __init__(self, ip, g_path='/home/ubuntu/f10_serverside/'):
        self.ip = ip
        self.ip_n = 0
        self.infoParamMap = {}
        # self.infoParamMap =pickle.load(open('categories.bin', 'rb'))
        self.interest = {'最新提示', '财务分析', '经营分析', '股东研究'}
        self.code_name = None

    def getMarket(self, c):
        c_m = {'600': 'sh', '601': 'sh', '603': 'sh', '000': 'sz', '001': 'sz', '002': 'sz', '300': 'sz'}
        m_n = {'sz': 0, 'sh': 1}
        for cm in c_m:
            if c.startswith(cm):
                return m_n[c_m[cm]]
        return None

    def setupCodeNameSina(self, g_path='/home/ubuntu/f10_serverside/'):

        import time

        import tushare as ts

        def get_today_df():
            try:
                return ts.get_today_all()
            except:
                time.sleep(60)
                return get_today_df()

        # code_name from today or the last trading day's data from Sina. Needless to
        # check the names again using the old way.
        # http://vip.stock.finance.sina.com.cn/quotes_service/api/json_v2.php/Market_Center.getHQNodeData?num=200&sort=code&asc=0&node=hs_a&symbol=&_s_r_a=page&page=1#
        def run(g_path):
            df = get_today_df()
            d = dict(zip(df.code, df.name))
            code_name = dict()
            for k in d:
                if k[:3] in {'600', '601', '603', '000', '001', '002', '300'}:
                    code_name[k] = d[k].replace(' ', '')

            with open(g_path + 'code_name.pkl', 'wb') as handle:
                pickle.dump(code_name, handle, protocol=pickle.HIGHEST_PROTOCOL)

            return code_name

        self.code_name = run(g_path)
        return len(self.code_name)

    def setupCodeName(self, max_retry=3, g_path='/home/ubuntu/f10_serverside/'):
        # get from all sources and pick the longest result
        code_name = []
        round_retry = 0
        api = TdxHq_API(auto_retry=True)
        while round_retry < max_retry:
            retry = 0
            ips = ["211.100.23.200:7779", "58.49.110.76:7709", "218.75.75.20:7709"]
            for addr in ips:
                h, p = addr.split(':')
                api_ret = api.connect(h, int(p))
                if not api_ret:
                    retry += 1
                    continue
                d = dict()
                for market in {0, 1}:
                    pos = 0
                    while True:
                        a = api.get_security_list(market, pos)
                        if a is None or len(a) == 0:
                            break
                        else:
                            pos += len(a)
                        for i in a:
                            code = i['code'].strip()
                            name = i['name'].replace(' ', '').strip('\x00')
                            if market == 0:
                                preset = {'000', '001', '002', '300'}
                            else:
                                preset = {'600', '601', '603'}
                            if code[:3] in preset:
                                d[code] = name
                    code_name.append(d)
                api.disconnect()
            if retry > 0:
                round_retry += 1
            else:
                break
        code_name.sort(key=len)
        self.code_name = code_name[-1]

        with open(g_path + 'code_name.pkl', 'wb') as handle:
            pickle.dump(self.code_name, handle, protocol=pickle.HIGHEST_PROTOCOL)

    def setupCategories(self, max_retry=3):
        import time
        api = None
        api = self._newApi(api)
        entry_count = 0
        for code in self.code_name:
            m = {}
            retry_count = 0
            while len(m) == 0 and retry_count < max_retry:
                raw = api.get_company_info_category(self.getMarket(code), code)
                if raw is not None:
                    for d in raw:
                        if d['name'] in self.interest:
                            entry_count += 1
                            m[d['name']] = (d['filename'], d['start'], d['length'])
                if len(m) == 0:
                    retry_count += 1
                    api = self._newApi(api)
                    time.sleep(1)
            if len(m) > 0:
                self.infoParamMap[code] = m
        print("got {} categories which should be x4 as #code on the market.".format(entry_count))
        self._closeApi(api)
        with open('categories.bin', 'wb') as handle:
            pickle.dump(self.infoParamMap, handle, protocol=pickle.HIGHEST_PROTOCOL)
        return entry_count

    def _closeApi(self, api):
        try:
            api.disconnect()
        except:
            pass

    def _newApi(self, api=None):
        if api is None:
            api = TdxHq_API(auto_retry=True, heartbeat=True)
        else:
            self._closeApi(api)
        host, port = self.ip[self.ip_n].split(':')
        print("using {} ...".format(self.ip[self.ip_n]))
        self.ip_n = (self.ip_n + 1) % len(self.ip)
        try:
            a = api.connect(host, int(port))
            return api if a else self._newApi(api)
        except Exception as e:
            print(str(e))
            return self._newApi(api)

    def getInfo(self, code_flag, flag_name_map, max_retry=3):

        res = []
        success = []
        retry_count = 0

        import time
        api = None
        while retry_count <= max_retry and len(code_flag) > 0:
            api = self._newApi(api)
            for (code, flag) in code_flag:
                flagname = flag_name_map[flag]
                try:
                    catTuple = self.infoParamMap[code][flagname]

                    raw = api.get_company_info_content(self.getMarket(code), code, catTuple[0], catTuple[1],
                                                       catTuple[2])
                    # it is VERY important to modify the source code to return the stream and IGNORE undecodable bytes
                    # can't override the original method
                    # '''
                    #     def parseResponse(self, body_buf):
                    #         pos = 0
                    #         _, length = struct.unpack(u'<10sH', body_buf[:12])
                    #         pos += 12
                    #         content = body_buf[pos: pos+length]
                    #         return content
                    # '''
                    raw_str = raw.decode('gbk', 'ignore') if len(raw) > 0 else ""

                    if (raw_str is not None) and (len(raw_str) > 0):
                        res.append({"code": code, "content": raw_str, "flag": flag})
                        success.append((code, flag))
                    else:

                        print("bytearray:{} EMPTY resp of {} in {}, retryCount {}.".format(len(raw), flagname, code,
                                                                                           retry_count))

                except Exception as e:
                    print("error occurred: {} in {}: {}".format(flagname, code, str(e)))

            for i in success:
                code_flag.remove(i)
            success.clear()
            retry_count += 1
            time.sleep(2)
        self._closeApi(api)
        for (code, flag) in code_flag:
            res.append({"code": code, "content": "EMPTY", "flag": flag})

        return res


if __name__ == '__main__':
    ip = ["211.100.23.200:7779", "58.49.110.76:7709", "101.71.255.135:7709", "211.100.23.202:7709", "218.75.75.20:7709"]
    hq = HQ(ip, "")
    print(hq.setupCodeNameSina(""))
    print(len(hq.code_name))
